
from src.util.depth_transform import DepthNormalizerBase, ScaleShiftDepthNormalizer
from src.util.alignment import align_depth_least_square
import numpy as np
import torch

gt_depth_path = 'sample_out/depth_npy/depth.npy'
pred_depth_path = 'sample_out/depth_npy/img_pred.npy'
# pred_depth_path = 'sample_out/depth_npy/depth_gt.npy'

gt_depth = torch.tensor(np.load(gt_depth_path))
pred_depth = torch.tensor(np.load(pred_depth_path))
valid_mask = torch.ones_like(torch.as_tensor(gt_depth)).bool()
normalizer = ScaleShiftDepthNormalizer(norm_min=-1.0,norm_max=1.0,min_max_quantile=0.02,clip=False)
norm_depth = normalizer(gt_depth)
# norm_depth = (norm_depth - pred_depth.min()) / (pred_depth.max() - pred_depth.min())
recovered_depth, scale, shift = align_depth_least_square(gt_arr=gt_depth,pred_arr=norm_depth,valid_mask_arr=valid_mask)
pred_depth = pred_depth * 2 - 1
pred_depth = pred_depth * scale + shift



# aligned_pred_depth,_,_ = align_depth_least_square(gt_arr=gt_depth,pred_arr=pred_depth,valid_mask_arr=valid_mask)


np.save('norm_test/gt_depth.npy',gt_depth.numpy())
np.save('norm_test/norm_depth.npy',norm_depth.numpy())
np.save('norm_test/recovered_depth.npy',recovered_depth.numpy())
np.save('norm_test/pred_depth.npy',pred_depth.numpy())

